#include "../Base.h"
#include <ntifs.h>


#ifndef _WIN64

NTSTATUS SetProcessHeaderFlags(PEPROCESS process)
{
	PUCHAR objectHeader;
#ifdef _WIN64
	objectHeader = ((PUCHAR)process - 0x30);

	*(objectHeader + 0x1b) |= 0x4;
#else
	objectHeader = ((PUCHAR)process - 0x18);

	*(objectHeader + 0xf) |= 0x4;
#endif

	return STATUS_SUCCESS;
}


//Short Name
PEPROCESS FindProcessByNameInBaseName(char* name)
{
	int i;
	PEPROCESS findedProcess = NULL;
	//�������̾��
	for (i = 0; i < 0x1000000; i += 4)
	{
		PEPROCESS process;
		PUCHAR	processName;

		NTSTATUS status = PsLookupProcessByProcessId((HANDLE)i, &process);
		if (!NT_SUCCESS(status))
		{
			continue;

		}
		processName = PsGetProcessImageFileName(process);
		if (processName && _stricmp((char*)processName, name) == 0)
		{
			findedProcess = process;
			break;
		}
		ObDereferenceObject(process);
	}
	return findedProcess;
}




//long name
PEPROCESS FindProcesByNameInFullName(WCHAR* name)
{
	//
	//	PEPROCESS findedProcess;
	//	for(int i = 8; i < 100000;i++)
	//	{
	//		PEPROCESS process;
	//		PUNICODE_STRING pathName = NULL;
	//		NTSTATUS status = PsLookupProcessByProcessId((HANDLE)i, &process);
	//		if(!NT_SUCCESS(status))
	//		{
	//			continue;
	//		}
	//	
	//		status = SeLocateProcessImageName(process, &pathName);
	//		if(!NT_SUCCESS(status))
	//		{
	//			ObDereferenceObject(process);
	//		}
	//		if(pathName->Length)
	//		{
	//			_wcsupr(pathName->Buffer);
	//			if (process && wcsstr(pathName->Buffer, name) != 0)
	//			{
	//				findedProcess = process;
	//				break;
	//			}
	//		}
	//		ExFreePoolWithTag(pathName, 0);
	//		ObDereferenceObject(process);
	//	}
	//	
	//	return findedProcess;
	//
	return NULL;
}


//��ȡ�������EPROCESS�е�ƫ��
UINT_PTR GetProcessActiveProcessLinksOffset()
{

	//static

	UNICODE_STRING funName;
	PUCHAR funAddress;
	int i;
	RtlInitUnicodeString(&funName, L"PsGetProcessId");
	funAddress = (PUCHAR)MmGetSystemRoutineAddress(&funName);
	if (funAddress)
	{
		for (i = 0; i < 100; i++)
		{
			if (funAddress[i] == 0x8B && funAddress[i + 1] == 0x80)
			{
				return *(ULONG*)(funAddress + i + 2) + 4;
			}
		}
	}
	return NULL;
}

void LogAllActiveProcessList()
{

	PEPROCESS process;
	UINT_PTR offset;
	PLIST_ENTRY list;
	PLIST_ENTRY nextEntry;
	PEPROCESS tempProcess;
	UCHAR* processName;
	process = PsGetCurrentProcess();
	if (!process)
	{
		return;
	}
	offset = GetProcessActiveProcessLinksOffset();
	list = (PLIST_ENTRY)((PUCHAR)process + offset);
	LOG("LIST address = %llx\r\n", (ULONG)list);
	LOG("LIST address = %llx\r\n", (ULONG)list->Blink);
	nextEntry = list->Blink;
	while (nextEntry && nextEntry != list)
	{

		tempProcess = (PEPROCESS)((PUCHAR)nextEntry - offset);
		processName = PsGetProcessImageFileName(tempProcess);
		LOG("[Process]    Name:%s\r\n", processName);
		nextEntry = nextEntry->Blink;
	}
}

VOID EnumThreadListOfProcess(char* name)
{
	PLIST_ENTRY nextEntry;
	LIST_ENTRY threadEntry;
	PEPROCESS_XP process;
	process = (PEPROCESS_XP)FindProcessByNameInBaseName(name);
	if (!process)
	{
		return;
	}
	threadEntry = process->ThreadListHead;
	nextEntry = threadEntry.Blink;

	LOG("nextEntry = %x threadEntry = %x \r\n", nextEntry, &threadEntry);
	nextEntry = nextEntry->Blink;
	LOG("nextEntry = %x threadEntry = %x \r\n", nextEntry, &threadEntry);
	LOG("isTrue = %d \r\n", (ULONG32)nextEntry == (ULONG32)&threadEntry);
	nextEntry = nextEntry->Blink;
	LOG("nextEntry = %x threadEntry = %x \r\n", nextEntry, &threadEntry);

	while (nextEntry && nextEntry != &threadEntry)
	{
		HANDLE threadParentPid;
		PETHREAD_XP thread_xp;
		thread_xp = CONTAINING_RECORD(nextEntry, ETHREAD_XP, ThreadListEntry);
		threadParentPid = thread_xp->Cid.UniqueProcess;
		LOG("nextEntry = %x threadEntry = %x \r\n", nextEntry, threadEntry);
		LOG("thread parent id = %x\r\n", threadParentPid);
		nextEntry = nextEntry->Blink;
	}
}

BOOLEAN ProtectProcessByName(char* name)
{
	PVOID table_entry;
	table_entry = GetTableEntryByProcessName(name);
	if (!table_entry) return FALSE;
	PEPROCESS_XP process = (PEPROCESS_XP)FindProcessByNameInBaseName(name);
	if (!process)return FALSE;
	memset(table_entry, 0, 8);
	process->UniqueProcessId = 0;
	return TRUE;
}


#else
PVOID obHandle;
HANDLE gProtectedId = 0;
VOID SetProtectedProcessId(HANDLE pid) {
	gProtectedId = pid;
}


OB_PREOP_CALLBACK_STATUS
PreCallback(
	_In_ PVOID RegistrationContext,
	_Inout_ POB_PRE_OPERATION_INFORMATION OperationInformation
) {

	PEPROCESS currentProcess = IoGetCurrentProcess();
	PEPROCESS targetProcess = (PEPROCESS)OperationInformation->Object;
	HANDLE currentPid = PsGetProcessId(currentProcess);
	HANDLE  targetPid = PsGetProcessId(targetProcess);
	if (currentPid == gProtectedId || targetProcess != gProtectedId) {
		return OB_PREOP_SUCCESS;
	}
	if (OperationInformation->Operation == OB_OPERATION_HANDLE_CREATE) {
		OperationInformation->Parameters->CreateHandleInformation.DesiredAccess = 0;
		OperationInformation->Parameters->CreateHandleInformation.OriginalDesiredAccess = 0;
	}
	else {
		OperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess = 0;
		OperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess = 0;
	}

}


VOID
PostCallback(
	_In_ PVOID RegistrationContext,
	_In_ POB_POST_OPERATION_INFORMATION OperationInformation
) {

	OperationInformation->Parameters->CreateHandleInformation.GrantedAccess = PROCESS_ALL_ACCESS;
	OperationInformation->Parameters->DuplicateHandleInformation.GrantedAccess = PROCESS_ALL_ACCESS;
}

NTSTATUS InitObRegister() {

	NTSTATUS status = STATUS_UNSUCCESSFUL;
	PDRIVER_OBJECT pDriver = GetDriverObjectByName(L"\\Driver\\WMIxWDM");
	if (!pDriver)return STATUS_UNSUCCESSFUL;

	ULONG64 jmpRax = searchCode("ntoskrnl.exe", ".text", "FFE1", 0);
	if (!jmpRax) return STATUS_UNSUCCESSFUL;

	PKLDR_DATA_TABLE_ENTRY ldr = (PKLDR_DATA_TABLE_ENTRY)pDriver->DriverSection;
	ldr->Flags |= 0x20;
	

	OB_OPERATION_REGISTRATION obOp = { 0 };
	obOp.ObjectType = PsProcessType;
	obOp.Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
	obOp.PreOperation = (POB_PRE_OPERATION_CALLBACK)jmpRax;
	obOp.PostOperation = PostCallback;
	OB_CALLBACK_REGISTRATION ob = { 0 };

	ob.Version = ObGetFilterVersion();
	ob.OperationRegistrationCount = 1;
	ob.RegistrationContext = PreCallback;
	ob.OperationRegistration = &obOp;
	UNICODE_STRING altitude;
	RtlInitUnicodeString(&altitude, L"322600");
	ob.Altitude = altitude;




	PUCHAR findFunc = NULL;
	//WIN7
	if (IsWin7())
	{
		PUCHAR func = (PUCHAR)ObRegisterCallbacks;
		for (int i = 0; i < 0x500; i++)
		{
			if (func[i] == 0x74 && func[i + 2] == 0xe8 && func[i + 7] == 0x3b && func[i + 8] == 0xc3)
			{
				LARGE_INTEGER larger;
				larger.QuadPart = (ULONG64)(func + i + 7);
				larger.LowPart += *(PULONG)(func + i + 3);
				findFunc = (PUCHAR)(larger.QuadPart);
				break;
			}
		}
	}
	else
	{
		PUCHAR func = (PUCHAR)ObRegisterCallbacks;
		for (int i = 0; i < 0x500; i++)
		{
			if (func[i] == 0xBA && func[i + 5] == 0xe8 && func[i + 10] == 0x85 && func[i + 11] == 0xc0)
			{
				LARGE_INTEGER larger;
				larger.QuadPart = (ULONG64)(func + i + 10);
				larger.LowPart += *(PULONG)(func + i + 6);
				findFunc = (PUCHAR)(larger.QuadPart);
				break;
			}
		}
	}



	if (findFunc)
	{
		PHYSICAL_ADDRESS phy = MmGetPhysicalAddress(findFunc);
		PVOID mem = MmMapIoSpace(phy, 10, MmNonCached);
		if (mem)
		{
			UCHAR bufCode[10] = { 0 };
			UCHAR patch[] = { 0xb0,0x1,0xc3 };
			memcpy(bufCode, mem, 10);
			memcpy(mem, patch, sizeof(patch));
			status = ObRegisterCallbacks(&ob, &obHandle);
			memcpy(mem, bufCode, 10);
		}
	}

	return status;

}


NTSTATUS ProtectProcessByPid(HANDLE pid)
{
	NTSTATUS status = STATUS_SUCCESS;
	if (!obHandle) {
		status = InitObRegister();
	}
	SetProtectedProcessId(pid);
	return status;
}

VOID UnProtectProcess()
{
	if (obHandle) {
		ObUnRegisterCallbacks(obHandle);
		gProtectedId = 0;
	}
}



PEPROCESS FindProcess(char* processName)
{
	PEPROCESS eprocess = NULL;
	KAPC_STATE kapc = { 0 };
	for (int i = 8; i < 0x10000; i += 4)
	{
		PEPROCESS tempProcess = NULL;
		NTSTATUS status = PsLookupProcessByProcessId((HANDLE)i, &tempProcess);
		if (NT_SUCCESS(status))
		{
			char* name = PsGetProcessImageFileName(tempProcess);
			if (name && _stricmp(name, processName) == 0)
			{
				eprocess = tempProcess;
				break;
			}
			ObDereferenceObject(tempProcess);

		}
	}

	return eprocess;
}

ULONG_PTR GetDirectoryTokenOffset() {
	return 0x123;
}



NTSTATUS RaiseToSystemPrivi(PEPROCESS process)
{

	static ULONG_PTR pSystem = NULL;
	if (pSystem == NULL)
	{
		NTSTATUS  Status;
		PEPROCESS  EProcess = NULL;
		Status = PsLookupProcessByProcessId((HANDLE)4, &EProcess);    //EPROCESS
		if (NT_SUCCESS(Status))
		{
			pSystem = (ULONG_PTR)EProcess + GetDirectoryTokenOffset();
			ObDereferenceObject(EProcess);
		}
	}


	UINT_PTR System = *(UINT_PTR*)pSystem & 0xfffffffffffffff0;

	*(UINT_PTR*)((UINT_PTR)process + GetDirectoryTokenOffset()) = System;
	return STATUS_SUCCESS;
}

NTSTATUS SetProcessCallbackTest(HANDLE pid, ULONG64 FuncAddress)
{
	PEPROCESS process = NULL;
	NTSTATUS status = PsLookupProcessByProcessId((HANDLE)pid, &process);
	if (!NT_SUCCESS(status)) {
		return status;
	}

	if (PsGetProcessExitStatus(process) != STATUS_PENDING) {
		ObDereferenceObject(process);
		return STATUS_PROCESS_IS_TERMINATING;
	}

	KAPC_STATE apc_state = { 0 };
	KeStackAttachProcess(process, &apc_state);
	PVOID BaseAddress = NULL;
	SIZE_T size = PAGE_SIZE;
	status = ZwAllocateVirtualMemory(NtCurrentProcess(), &BaseAddress, 0, &size, MEM_COMMIT, PAGE_EXECUTE_READWRITE);

	memset(BaseAddress, 0, size);
	DbgBreakPoint();
	char bufCode[] =
	{ 0xcc,
					0x50,
					0x51,
					0x52,
					0x53,
					0x55,
					0x56,
					0x57,
					0x41,0x50,
					0x41,0x51,
					0x41,0x52,
					0x41,0x53,
					0x41,0x54,
					0x41,0x55,
					0x41,0x56,
					0x41,0x57,
					0x48,0xB9,0x21,0x21,0x21,0x21,0x01,0x00,0x00,0x00,
					0x49,0x39,0xcA,
		//0xB9,0x05,0x00,0x00,0xC0,
		0x41,0x5F,
		0x41,0x5E,
		0x41,0x5D,
		0x41,0x5C,
		0x41,0x5B,
		0x41,0x5A,
		0x41,0x59,
		0x41,0x58,
		0x5F,
		0x5E,
		0x5D,
		0x5B,
		0x5A,
		0x59,
		0x41,0xFF,0xE2
	};
	*(PULONG64)&bufCode[26] = (ULONG64)NtOpenProcess;




	memcpy(BaseAddress, bufCode, sizeof(bufCode));
	KeUnstackDetachProcess(&apc_state);
	SetProcessCallbackDo((HANDLE)2596, (ULONG64)BaseAddress);
	ObDereferenceObject(process);
	return status;
}



PEPROCESS GetProcessById(HANDLE pid)
{

	PEPROCESS process;
	NTSTATUS status = PsLookupProcessByProcessId(pid, &process);
	if (!NT_SUCCESS(status)) {
		return NULL;
	}

	if (PsGetProcessExitStatus(process) != STATUS_PENDING) {
		ObDereferenceObject(process);
		return NULL;
	}

	return process;
}

 


//Short Name
PEPROCESS FindProcessByNameInBaseName(char* name)
{
	int i;
	PEPROCESS findedProcess = NULL;
	//�������̾��
	for (i = 0; i < 0x1000000; i += 4)
	{
		PEPROCESS process;
		PUCHAR	processName;

		NTSTATUS status = PsLookupProcessByProcessId((HANDLE)i, &process);
		if (!NT_SUCCESS(status))
		{
			continue;

		}
		processName = (PUCHAR)PsGetProcessImageFileName(process);
		if (processName && _stricmp((char*)processName, name) == 0)
		{
			findedProcess = process;
			break;
		}
		ObDereferenceObject(process);
	}
	return findedProcess;
}






NTSTATUS SetProcessCallbackDo(HANDLE pid, ULONG64 FuncAddressR3) {

	PEPROCESS process = NULL;
	NTSTATUS status = PsLookupProcessByProcessId(pid, &process);
	if (!NT_SUCCESS(status)) {
		return status;
	}

	if (PsGetProcessExitStatus(process) != STATUS_PENDING) {
		return STATUS_PROCESS_IS_TERMINATING;
	}

	KAPC_STATE apc_state = { 0 };
	KeStackAttachProcess(process, &apc_state);

	ULONG64 funAddr = FuncAddressR3;

	PACCESS_TOKEN token = PsReferencePrimaryToken(process);
	PULONG64 tokenAccess = (PULONG64)((PUCHAR)token + 0x40);
	tokenAccess[1] |= 0x100000;
	tokenAccess[2] |= 0x100000;
	tokenAccess[3] |= 0x100000;
	status = ZwSetInformationProcess(PsGetCurrentProcessId(), ProcessInstrumentationCallback, &funAddr, sizeof(ULONG64));
	KeUnstackDetachProcess(&apc_state);

	return status;
}




#endif
